import cv2
import numpy as np

def calculate_slope(line):

    x_1, y_1, x_2, y_2 = line[0]
    return (y_2-y_1)/(x_2-x_1)

edge_img = cv2.imread('cmh1.jpg',cv2.IMREAD_GRAYSCALE)

lines = cv2.HoughLinesP(edge_img,1,np.pi/180,15,minLineLength=40,maxLineGap=20)

left_lines = [line for line in lines if calculate_slope(line) > 0]
right_lines = [line for line in lines if calculate_slope(line) < 0]

def reject_abnormal_lines(lines,threshold):
    slopes = [calculate_slope(line) for line in lines]
    while len(lines) > 0:
        mean = np.mean(slopes)
        diff = [abs(s-mean) for s in slopes]
        idx = np.argmax(diff)
        if diff[idx]>threshold:
            slopes.pop(idx)
            lines.pop(idx)
        else:
            break
            return lines



    reject_abnormal_lines(left_lines, threshold=0.2)
    reject_abnormal_lines(right_lines, threshold=0.2)

def least_squares_fit(lines):

 x_coords = np.ravel([[line[0][0],line[0][2]] for line in lines])
 y_coords = np.ravel([[line[0][1],line[0][3]] for line in lines])

 poly = np.polyfit(x_coords,y_coords,deg=1)

 point_min = (np.min(x_coords),np.polyval(poly,np.min(x_coords)))
 point_max = (np.max(x_coords),np.polyval(poly,np.max(x_coords)))

 return np.array([point_min,point_max],dtype=np.int)


left_lines=least_squares_fit(left_lines)

right_lines=least_squares_fit(right_lines)

img = cv2.imread('hxy.jpg',cv2.IMREAD_GRAYSCALE)
cv2.line(img, tuple(left_lines[0]),tuple(left_lines[1]),color=(0,255,255),thickness=5)
cv2.line(img, tuple(right_lines[0]),tuple(right_lines[1]),color=(0,255,255),thickness=5)
cv2.imshow('lane',img)
cv2.waitKey(0)